import json
import logging
import operator
from typing import TypedDict, List, Annotated

from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.tools import render_text_description
from langgraph.graph import END, StateGraph
from langgraph.graph.state import CompiledStateGraph

from core.base import chat_llm
from graph_agent.self.parse import stateParse, msgParse
from graph_agent.self.prompts import promptTemplate
from util.tools import get_weather, searxng_search, milvus_search

logger = logging.getLogger("graph")



# 初始化模型
llm = chat_llm()

prompt = ChatPromptTemplate.from_messages([
    ("system","你是非常强大的助手，你可以使用各种工具来完成人类交给的问题和任务"),
    MessagesPlaceholder("chat_history"),
    ("user",promptTemplate),
    MessagesPlaceholder(variable_name="agent_scratchpad")
])

tools = [get_weather,searxng_search,milvus_search]






class GState(TypedDict):
    messages: Annotated[List, operator.add]
    chat_history: Annotated[List, operator.add]

def isUseTool(gState):
    obj = gState["messages"][-1]
    if obj["action"] == "Final Answer":
        return "no"
    else:
        return "yes"


def useTool(gState):
    obj = gState["messages"][-1]
    for tool  in tools:
        if obj["action"] == tool.name:
            # print(obj["action_input"])
            tool_answer = tool.invoke(json.loads(obj["action_input"]))
            print(f"tool_answer:{tool_answer}" )
            return {"messages":[tool_answer],"action":'tool:'+tool.name}
    print("找不到对应的工具名称：",obj["action"])
    return None

async def create_agent() -> CompiledStateGraph:
    pormpt = prompt.partial(
        tools=render_text_description(tools),
        tool_names=", ".join([t.name for t in tools])
    )

    chain = pormpt | llm
    agentGraph = StateGraph(GState)

    agentGraph.add_node("startNode", stateParse | chain | msgParse)

    agentGraph.add_node("toolNode",  useTool)

    agentGraph.add_conditional_edges("startNode", isUseTool, {"no": END, "yes": "toolNode"})

    agentGraph.add_edge("toolNode", "startNode")

    agentGraph.set_entry_point("startNode")
    # try:
    #     agent.get_graph().draw_mermaid_png(output_file_path="graph.png")
    # except Exception:
    #     # This requires some extra dependencies and is optional
    #     pass
    # agent.invoke({"messages":["刘亦菲最近有什么活动？"]})
    return agentGraph.compile()



